import os
from datetime import datetime
import matplotlib
from matplotlib import pyplot as PLT
from scipy.stats import linregress
import numpy

def Polyfnc(x, m, b):
    return m*x + b

def PlotMaxtvsVel(database, bins, max_d):
    
    max_torque = []
    max_vel = []
    max_at_B2 = []
    max_at_B3 = []
    filelist = []
    avgtor = []
    avgvel = []
    linfit = []
    linfit_B2 = []
    linfit_B3 = []
    linfit2 = []
    linfit3 = []
    
    for run in database:
        if run.Direction() != 'in': continue
        if run.Source() != 'weakTh-228': continue
        if run.Station() != 'S5': continue
        if max(run.NewPosition()) < 350: continue   # because some stations in run_history are wrong
        if max(run.NewPosition()) > 400: continue
        print run
        
        #print len(data['rpm']), len(data['velocity'])
        #for i, item in enumerate(data['rpmvel']):
        #    print data['RPM'][i], '         ', data['rpm'][i], data['velocity'][i]
        
        filelist.append(run.Filename().split('.')[0].split(' ')[2].split('_')[0])
        max_torque.append(run.MaxTorque())
        max_vel.append(max(run.Velocity()))
        
        X = numpy.linspace(0,max_d,bins)
        T = numpy.zeros(bins)
        V = numpy.zeros(bins)
        
        # bin
        for j,x in enumerate(X):
            pos_torque = []
            pos_vel = []
            x0 = x-(max_d/bins)/2.
            x1 = x+(max_d/bins)/2.
            for i,p in enumerate(run.Position()):
                if p>x0 and p<=x1:
                    pos_torque.append(run.Torque()[i])
            if len(pos_torque)>0:
                T[j] += max(pos_torque)
        
        # average torque
        t = 0
        for torque in T:
            t += torque
        t2 = t / len(T)
        avgtor.append(t2)
        
        # average velocity
        sumvel = 0
        for j, v in enumerate(run.Velocity()):
            sumvel += v
        averagevel = sumvel / int(len(run.Velocity()))
        avgvel.append(averagevel)
       
        pos_torque = []
        for i, x in enumerate(run.NewPosition()):
            if x > 300 and x < 330:
                pos_torque.append(run.Torque()[i])
        max_at_B2.append(max(pos_torque))
        
        pos_torque2 = []
        for i, x in enumerate(run.NewPosition()):
            if x > 350 and x < 375:
                pos_torque2.append(run.Torque[i])
        max_at_B3.append(max(pos_torque2))
        
    
    lin = linregress(avgvel, max_torque)
    lin_B2 = linregress(avgvel, max_at_B2)
    lin_B3 = linregress(avgvel, max_at_B3)
    lin2 = linregress(avgvel, avgtor)
    lin3 = linregress(max_vel, max_torque)
    
    for thing in avgvel:
        linfit.append(Polyfnc(thing, lin[0], lin[1]))
        linfit_B2.append(Polyfnc(thing, lin_B2[0], lin_B2[1]))
        linfit_B3.append(Polyfnc(thing, lin_B3[0], lin_B3[1]))
        linfit2.append(Polyfnc(thing, lin2[0], lin2[1]))
    for other in max_vel:
        linfit3.append(Polyfnc(other, lin3[0], lin3[1]))
    
    # plot
    box = [0.14, 0.14, 0.76, 0.76]
    
    fig = PLT.figure(figsize=(15,8),dpi=150)
    ax1 = fig.add_axes(box)
    ax1.set_ylabel('Max Torque (in.oz)')
    ax1.set_xlabel('Average Velocity (cm/s)')
    PLT.title('Max Torque Versus Average Velocity')
    ax1.scatter(avgvel, max_torque, color = 'g')
    ax1.plot(avgvel, linfit, '-', color = 'g')
    #for i, txt in enumerate(filelist):
        #ax1.annotate(txt, (date[i], max_torque[i]), position = (date[i], max_torque[i]), size = 'x-small', alpha = 0.5, rotation = 45)
    fig.savefig('vel/maxt_vs_vel.png')
    
    fig2 = PLT.figure(figsize=(15,8),dpi=150)
    ax2 = fig2.add_axes(box)
    ax2.set_ylabel('Max Torque (in.oz)')
    ax2.set_xlabel('Date')
    PLT.title('Max Torque Versus Average Velocity for Bend A6 and A7')
    ax2.scatter(avgvel, max_at_B2, color = 'b', label = 'Bend A6')
    ax2.plot(avgvel, linfit_B2, '-', color = 'b')
    ax2.scatter(avgvel, max_at_B3, color = 'r', marker = 'D', label = 'Bend A7')
    ax2.plot(avgvel, linfit_B3, '-', color = 'r')
    fig2.savefig('vel/maxt_vs_vel_for_A67.png')
    
    fig4 = PLT.figure(figsize=(15,8),dpi=150)
    ax4 = fig4.add_axes(box)
    ax4.set_ylabel('Average Torque (in.oz)')
    ax4.set_xlabel('Average Velocity (cm/s)')
    PLT.title('Average Torque Versus Average Velocity')
    ax4.scatter(avgvel, avgtor, color = 'b')
    ax4.plot(avgvel, linfit2, '-', color = 'b')
    fig4.savefig('vel/avgt_vs_vel.png')
    
    fig5 = PLT.figure(figsize=(15,8),dpi=150)
    ax5 = fig5.add_axes(box)
    ax5.set_ylabel('Maximum Torque (in.oz)')
    ax5.set_xlabel('Maximum Velocity (cm/s)')
    PLT.title('Maximum Torque Versus Maximum Velocity')
    ax5.scatter(max_vel, max_torque, color = 'g')
    fig5.savefig('vel/maxt_vs_maxv.png')
    